{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import re\n",
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "from random import *\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as Data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 这里并没有用什么大型的数据集，而是手动输入了两个人的对话，\n",
    "# 主要是为了降低代码阅读难度，我希望读者能更关注模型实现的部分\n",
    "text = (\n",
    "    'Hello, how are you? I am Romeo.\\n'                     # R #   你好，你过得怎么样？我是罗密欧。    \n",
    "    'Hello, Romeo My name is Juliet. Nice to meet you.\\n'   # J # 你好，罗密欧，你叫什么名字？很高兴认识你。\n",
    "    'Nice meet you too. How are you today?\\n'               # R # 很高兴认识你，你今天过得怎么样？\n",
    "    'I am doing well. How about you?\\n'                     # J # 我很好，你呢？\n",
    "    'I am doing great. What about you?\\n'                   # R # 我很好，你呢？\n",
    "    'Great. My baseball team won the competition.\\n'        # J # 很好，我的棒球队赢得了比赛。\n",
    "    'Congratulations, Juliet\\n'                             # R # 太棒了，朱丽叶。\n",
    "    'Thank you Romeo\\n'                                     # J # 谢谢罗密欧。\n",
    "    'Where are you going today?\\n'                          # R # 你今天要去哪里？\n",
    "    'I am going to visit my grandmother. \\n'                # J # 我今天要去看我奶奶。\n",
    ")\n",
    "sentences = re.sub(\"[.,!?\\\\-]\", '', text.lower()).split('\\n') # filter '.', ',', '?', '!'\n",
    "word_list = list(set(\" \".join(sentences).split())) # ['hello', 'how', 'are', 'you',...]\n",
    "word2idx = {'[PAD]' : 0, '[CLS]' : 1, '[SEP]' : 2, '[MASK]' : 3}\n",
    "for i, w in enumerate(word_list):\n",
    "    word2idx[w] = i + 4\n",
    "idx2word = {i: w for i, w in enumerate(word2idx)}\n",
    "vocab_size = len(word2idx)\n",
    "\n",
    "token_list = list()\n",
    "for sentence in sentences:\n",
    "    arr = [word2idx[s] for s in sentence.split()]\n",
    "    token_list.append(arr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#最终token_list是个二维的list，里面每一行代表一句话\n",
    "print(token_list)\n",
    "'''\n",
    "[[12, 7, 22, 5, 39, 21, 15],\n",
    " [12, 15, 13, 35, 10, 27, 34, 14, 19, 5],\n",
    " [34, 19, 5, 17, 7, 22, 5, 8],\n",
    " [33, 13, 37, 32, 28, 11, 16],\n",
    " [30, 23, 27],\n",
    " [6, 5, 15],\n",
    " [36, 22, 5, 31, 8],\n",
    " [39, 21, 31, 18, 9, 20, 5],\n",
    " [39, 21, 31, 14, 29, 13, 4, 25, 10, 26, 38, 24]]\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#模型超参数\n",
    "# BERT Parameters\n",
    "# n_segments表示Decoder input由几句话组成\n",
    "maxlen = 30         # 输入句子的最大长度\n",
    "batch_size = 6      # 批处理尺寸\n",
    "max_pred = 5        # 最大的预测词数\n",
    "n_layers = 6        # 编码器的层数\n",
    "n_heads = 12         # 多头注意力的头数\n",
    "d_model = 768        # 词向量维度\n",
    "d_ff = d_model*4     # 4 倍的模型维度，前馈网络维度  \n",
    "d_k = d_v = 64       # 多头注意力机制 d_k：d_v它是键（Key）向量的维度\n",
    "n_segments = 2       # Decoder input由几句话组成\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据预处理\n",
    "# 数据预处理部分，我们需要根据概率随机make或者替换\n",
    "# （以下统称mask）一句话中15%的token，还需要拼接任意两句话\n",
    "\n",
    "# sample IsNext and NotNext to be same in small batch size\n",
    "def make_data():\n",
    "    batch = []\n",
    "    positive = negative = 0\n",
    "    while positive != batch_size/2 or negative != batch_size/2:\n",
    "        # sample random index in sentences\n",
    "        tokens_a_index, tokens_b_index = randrange(len(sentences)), randrange(len(sentences)) \n",
    "        tokens_a, tokens_b = token_list[tokens_a_index], token_list[tokens_b_index]\n",
    "        input_ids = [word2idx['[CLS]']] + tokens_a + [word2idx['[SEP]']] + tokens_b + [word2idx['[SEP]']]\n",
    "        segment_ids = [0] * (1 + len(tokens_a) + 1) + [1] * (len(tokens_b) + 1)\n",
    "\n",
    "        # MASK LM\n",
    "         # 15 % of tokens in one sentence\n",
    "        n_pred =  min(max_pred, max(1, int(len(input_ids) * 0.15)))\n",
    "        # candidate masked position\n",
    "        cand_maked_pos = [i for i, token in enumerate(input_ids)\n",
    "                          if token != word2idx['[CLS]'] and token != word2idx['[SEP]']] \n",
    "        shuffle(cand_maked_pos)\n",
    "        masked_tokens, masked_pos = [], []\n",
    "        for pos in cand_maked_pos[:n_pred]:\n",
    "            masked_pos.append(pos)\n",
    "            masked_tokens.append(input_ids[pos])\n",
    "            if random() < 0.8:  # 80%\n",
    "                input_ids[pos] = word2idx['[MASK]'] # make mask\n",
    "            elif random() > 0.9:  # 10%\n",
    "                index = randint(0, vocab_size - 1) # random index in vocabulary\n",
    "                while index < 4: # can't involve 'CLS', 'SEP', 'PAD'\n",
    "                  index = randint(0, vocab_size - 1)\n",
    "                input_ids[pos] = index # replace\n",
    "\n",
    "        # Zero Paddings\n",
    "        n_pad = maxlen - len(input_ids)\n",
    "        input_ids.extend([0] * n_pad)\n",
    "        segment_ids.extend([0] * n_pad)\n",
    "\n",
    "        # Zero Padding (100% - 15%) tokens\n",
    "        if max_pred > n_pred:\n",
    "            n_pad = max_pred - n_pred\n",
    "            masked_tokens.extend([0] * n_pad)\n",
    "            masked_pos.extend([0] * n_pad)\n",
    "\n",
    "        if tokens_a_index + 1 == tokens_b_index and positive < batch_size/2:\n",
    "            batch.append([input_ids, segment_ids, masked_tokens, masked_pos, True]) # IsNext\n",
    "            positive += 1\n",
    "        elif tokens_a_index + 1 != tokens_b_index and negative < batch_size/2:\n",
    "            batch.append([input_ids, segment_ids, masked_tokens, masked_pos, False]) # NotNext\n",
    "            negative += 1\n",
    "    return batch\n",
    "# Proprecessing Finished\n",
    "\n",
    "batch = make_data()\n",
    "input_ids, segment_ids, masked_tokens, masked_pos, isNext = zip(*batch)\n",
    "input_ids, segment_ids, masked_tokens, masked_pos, isNext = \\\n",
    "    torch.LongTensor(input_ids),  torch.LongTensor(segment_ids), torch.LongTensor(masked_tokens),\\\n",
    "    torch.LongTensor(masked_pos), torch.LongTensor(isNext)\n",
    "\n",
    "class MyDataSet(Data.Dataset):\n",
    "  def __init__(self, input_ids, segment_ids, masked_tokens, masked_pos, isNext):\n",
    "    self.input_ids = input_ids\n",
    "    self.segment_ids = segment_ids\n",
    "    self.masked_tokens = masked_tokens\n",
    "    self.masked_pos = masked_pos\n",
    "    self.isNext = isNext\n",
    "  \n",
    "  def __len__(self):\n",
    "    return len(self.input_ids)\n",
    "  \n",
    "  def __getitem__(self, idx):\n",
    "    return self.input_ids[idx], self.segment_ids[idx], self.masked_tokens[idx], self.masked_pos[idx], self.isNext[idx]\n",
    "\n",
    "loader = Data.DataLoader(MyDataSet(input_ids, segment_ids, masked_tokens, masked_pos, isNext), batch_size, True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#模型构建\n",
    "#模型结构主要采用了Transformer的Encoder\n",
    "def get_attn_pad_mask(seq_q, seq_k):\n",
    "    batch_size, seq_len_q = seq_q.size()\n",
    "    batch_size, seq_len_k = seq_k.size()\n",
    "    # 这里假设 PAD 的索引为 0，根据实际情况修改\n",
    "    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size, 1, seq_len_k]\n",
    "    return pad_attn_mask.expand(batch_size, seq_len_q, seq_len_k)  # [batch_size, seq_len_q, seq_len_k]\n",
    "def gelu(x):\n",
    "    \"\"\"\n",
    "    GELU（高斯误差线性单元）激活函数的实现。\n",
    "    相关信息：OpenAI GPT 所使用的 GELU 激活函数略有不同（且计算结果也会稍有差异）：\n",
    "      0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))\n",
    "      Also see https://arxiv.org/abs/1606.08415\n",
    "    \"\"\"\n",
    "    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))\n",
    "\n",
    "class Embedding(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Embedding, self).__init__()\n",
    "        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding\n",
    "        self.pos_embed = nn.Embedding(maxlen, d_model)  # position embedding\n",
    "        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding\n",
    "        self.norm = nn.LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, x, seg):\n",
    "        seq_len = x.size(1)\n",
    "        pos = torch.arange(seq_len, dtype=torch.long)\n",
    "        pos = pos.unsqueeze(0).expand_as(x)  # [seq_len] -> [batch_size, seq_len]\n",
    "        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)\n",
    "        return self.norm(embedding)\n",
    "\n",
    "class ScaledDotProductAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ScaledDotProductAttention, self).__init__()\n",
    "\n",
    "    def forward(self, Q, K, V, attn_mask):\n",
    "        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # scores : [batch_size, n_heads, seq_len, seq_len]\n",
    "        scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one.\n",
    "        attn = nn.Softmax(dim=-1)(scores)\n",
    "        context = torch.matmul(attn, V)\n",
    "        return context\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.W_Q = nn.Linear(d_model, d_k * n_heads)\n",
    "        self.W_K = nn.Linear(d_model, d_k * n_heads)\n",
    "        self.W_V = nn.Linear(d_model, d_v * n_heads)\n",
    "    def forward(self, Q, K, V, attn_mask):\n",
    "        # q: [batch_size, seq_len, d_model], k: [batch_size, seq_len, d_model], v: [batch_size, seq_len, d_model]\n",
    "        residual, batch_size = Q, Q.size(0)\n",
    "        # (B, S, D) -proj-> (B, S, D) -split-> (B, S, H, W) -trans-> (B, H, S, W)\n",
    "        q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # q_s: [batch_size, n_heads, seq_len, d_k]\n",
    "        k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # k_s: [batch_size, n_heads, seq_len, d_k]\n",
    "        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2)  # v_s: [batch_size, n_heads, seq_len, d_v]\n",
    "\n",
    "        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # attn_mask : [batch_size, n_heads, seq_len, seq_len]\n",
    "\n",
    "        # context: [batch_size, n_heads, seq_len, d_v], attn: [batch_size, n_heads, seq_len, seq_len]\n",
    "        context = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)\n",
    "        context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v) # context: [batch_size, seq_len, n_heads, d_v]\n",
    "        output = nn.Linear(n_heads * d_v, d_model)(context)\n",
    "        return nn.LayerNorm(d_model)(output + residual) # output: [batch_size, seq_len, d_model]\n",
    "\n",
    "class PoswiseFeedForwardNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PoswiseFeedForwardNet, self).__init__()\n",
    "        self.fc1 = nn.Linear(d_model, d_ff)\n",
    "        self.fc2 = nn.Linear(d_ff, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_ff) -> (batch_size, seq_len, d_model)\n",
    "        return self.fc2(gelu(self.fc1(x)))\n",
    "\n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.enc_self_attn = MultiHeadAttention()\n",
    "        self.pos_ffn = PoswiseFeedForwardNet()\n",
    "\n",
    "    def forward(self, enc_inputs, enc_self_attn_mask):\n",
    "        enc_outputs = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_inputs to same Q,K,V\n",
    "        enc_outputs = self.pos_ffn(enc_outputs) # enc_outputs: [batch_size, seq_len, d_model]\n",
    "        return enc_outputs\n",
    "\n",
    "class BERT(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BERT, self).__init__()\n",
    "        self.embedding = Embedding()\n",
    "        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(d_model, d_model),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        self.classifier = nn.Linear(d_model, 2)\n",
    "        self.linear = nn.Linear(d_model, d_model)\n",
    "        self.activ2 = gelu\n",
    "        # fc2 is shared with embedding layer\n",
    "        embed_weight = self.embedding.tok_embed.weight\n",
    "        self.fc2 = nn.Linear(d_model, vocab_size, bias=False)\n",
    "        self.fc2.weight = embed_weight\n",
    "\n",
    "    def forward(self, input_ids, segment_ids, masked_pos):\n",
    "        output = self.embedding(input_ids, segment_ids) # [bach_size, seq_len, d_model]\n",
    "        enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids) # [batch_size, maxlen, maxlen]\n",
    "        for layer in self.layers:\n",
    "            # output: [batch_size, max_len, d_model]\n",
    "            output = layer(output, enc_self_attn_mask)\n",
    "        # it will be decided by first token(CLS)\n",
    "        h_pooled = self.fc(output[:, 0]) # [batch_size, d_model]\n",
    "        logits_clsf = self.classifier(h_pooled) # [batch_size, 2] predict isNext\n",
    "\n",
    "        masked_pos = masked_pos[:, :, None].expand(-1, -1, d_model) # [batch_size, max_pred, d_model]\n",
    "        h_masked = torch.gather(output, 1, masked_pos) # masking position [batch_size, max_pred, d_model]\n",
    "        h_masked = self.activ2(self.linear(h_masked)) # [batch_size, max_pred, d_model]\n",
    "        logits_lm = self.fc2(h_masked) # [batch_size, max_pred, vocab_size]\n",
    "        return logits_lm, logits_clsf\n",
    "model = BERT()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adadelta(model.parameters(), lr=0.001)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 用到了torch.gather()函数，这个函数实际上实现了以下的功能\n",
    "#out = torch.gather(input, dim, index)\n",
    "# out[i][j][k] = input[index[i][j][k]][j][k] # dim=0\n",
    "# out[i][j][k] = input[i][index[i][j][k]][k] # dim=1\n",
    "# out[i][j][k] = input[i][j][index[i][j][k]] # dim=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 具体以一个例子来说就是，首先我生成index变量\n",
    "index = torch.from_numpy(np.array([[1, 2, 0], [2, 0, 1]])).type(torch.LongTensor)\n",
    "index = index[:, :, None].expand(-1, -1, 10)\n",
    "print(index)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 然后随机生成一个[2, 3, 10]维的tensor，可以理解为有2个batch\n",
    "# 每个batch有3句话，每句话由10个词构成\n",
    "# 只不过这里的词不是以正整数（索引）的形式出现，而是连续的数值\n",
    "\n",
    "input = torch.rand(2, 3, 10)\n",
    "print(input)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0060 loss = 0.785403\n",
      "Epoch: 0070 loss = 0.813557\n",
      "Epoch: 0080 loss = 0.739547\n",
      "Epoch: 0090 loss = 0.756417\n",
      "Epoch: 0100 loss = 0.732498\n",
      "Epoch: 0110 loss = 0.745188\n",
      "Epoch: 0120 loss = 0.770443\n",
      "Epoch: 0130 loss = 0.795457\n",
      "Epoch: 0140 loss = 0.761338\n",
      "Epoch: 0150 loss = 0.788986\n",
      "Epoch: 0160 loss = 0.782730\n",
      "Epoch: 0170 loss = 0.825565\n",
      "Epoch: 0180 loss = 0.744066\n"
     ]
    }
   ],
   "source": [
    "# 训练\n",
    "model.train()\n",
    "for epoch in range(180):\n",
    "    for input_ids, segment_ids, masked_tokens, masked_pos, isNext in loader:\n",
    "      logits_lm, logits_clsf = model(input_ids, segment_ids, masked_pos)\n",
    "      loss_lm = criterion(logits_lm.view(-1, vocab_size), masked_tokens.view(-1)) # for masked LM\n",
    "      loss_lm = (loss_lm.float()).mean()\n",
    "      loss_clsf = criterion(logits_clsf, isNext) # for sentence classification\n",
    "      loss = loss_lm + loss_clsf\n",
    "      if (epoch + 1) % 10 == 0:\n",
    "          print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))\n",
    "      optimizer.zero_grad()\n",
    "      loss.backward()\n",
    "      optimizer.step()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "文本内容： Hello, how are you? I am Romeo.\n",
      "Hello, Romeo My name is Juliet. Nice to meet you.\n",
      "Nice meet you too. How are you today?\n",
      "I am doing well. How about you?\n",
      "I am doing great. What about you?\n",
      "Great. My baseball team won the competition.\n",
      "Congratulations, Juliet\n",
      "Thank you Romeo\n",
      "Where are you going today?\n",
      "I am going to visit my grandmother. \n",
      "\n",
      "输入的词列表（去除填充符）： ['[CLS]', 'congratulations', '[MASK]', '[SEP]', 'i', 'am', 'going', 'to', 'visit', 'my', 'grandmother', '[SEP]']\n",
      "真实的掩码标记列表： [30]\n",
      "预测的掩码标记列表： [26]\n",
      "真实的句子相邻情况： 否\n",
      "预测的句子相邻情况： 否\n"
     ]
    }
   ],
   "source": [
    "# 测试\n",
    "# 预测掩码标记和句子是否相邻\n",
    "input_ids, segment_ids, masked_tokens, masked_pos, isNext = batch[0]\n",
    "print(\"文本内容：\", text)\n",
    "print(\"输入的词列表（去除填充符）：\", [idx2word[w] for w in input_ids if idx2word[w] != '[PAD]'])\n",
    "\n",
    "logits_lm, logits_clsf = model(torch.LongTensor([input_ids]),\n",
    "                               torch.LongTensor([segment_ids]), torch.LongTensor([masked_pos]))\n",
    "logits_lm = logits_lm.data.max(2)[1][0].data.numpy()\n",
    "print(\"真实的掩码标记列表：\", [pos for pos in masked_tokens if pos != 0])\n",
    "print(\"预测的掩码标记列表：\", [pos for pos in logits_lm if pos != 0])\n",
    "\n",
    "logits_clsf = logits_clsf.data.max(1)[1].data.numpy()[0]\n",
    "print(\"真实的句子相邻情况：\", \"是\" if isNext else \"否\")\n",
    "print(\"预测的句子相邻情况：\", \"是\" if logits_clsf else \"否\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
